#!/usr/bin/python
# -*- coding: utf-8 -*-

import Ice, IceGrid
import os, sys, time

try:
    HOME = os.environ['FIPA']
except KeyError:
    print 'FIPA environment variable is not currently defined.'
    sys.exit(0)

SLICE = os.path.join(HOME, 'slice')
FILE = os.path.join(HOME, 'test', 'Spammer.ice')
Ice.loadSlice('--all -I/usr/share/Ice/slice', ['-I' + SLICE, FILE])

import FIPA
import Test

class AMI_Model_beginTestI (object):

    def ice_response (self):
        pass

    def ice_exception (self, ex):
        print 'Client.py: ' + str(ex)

class Client (Ice.Application):

    def run (self, argv):

        if len(sys.argv) < 2:
            print 'Usage: python Client.py --create <number of agent pairs>'
            print '\t or: python Client.py --info'
            print '\t or: python Client.py --begin <spam messages> <message size>'
            print '\t or: python Client.py --begin-with-mts <spam messages> <message size>'
            print '\t or: python Client.py --clean'
        else:

            if sys.argv[1] == '--create':
                self.createUserSpammerPairs(int(sys.argv[2]), self.communicator())
            if sys.argv[1] == '--info':
                self.infoUserSpammerPairs(self.communicator())
            if sys.argv[1] == '--begin':
                self.beginTest(int(sys.argv[2]), False, int(sys.argv[3]), self.communicator())
            if sys.argv[1] == '--begin-with-mts':
                self.beginTest(int(sys.argv[2]), True, int(sys.argv[3]), self.communicator())
            if sys.argv[1] == '--clean':
                self.clean(self.communicator())

        return 0

    def createUserSpammerPairs (self, number, comm):

        query = IceGrid.QueryPrx.checkedCast(comm.stringToProxy('FIPA/Query'))
        df = FIPA.DFPrx.checkedCast(query.findObjectById(comm.stringToIdentity('DF')))

        # Get all the agent factories.
        factories = query.findAllObjectsByType('::FIPA::AgentFactory')

        users = []

        print 'Creating the user agents...'
        for f in factories:
            fact = FIPA.AgentFactoryPrx.checkedCast(f)
            for i in range (0, number):
                users.append(fact.create(''))

        for user in users:
            df.modify(FIPA.DfAgentDescription(user, [FIPA.ServiceDescription('user')]))

        # Get all the spammer factories.
        factories = query.findAllObjectsByType('::Test::SpammerFactory')

        # Create the spammer agents.
        spammers = []

        print 'Creating the spammer agents...'
        for f in factories:
            fact = FIPA.AgentFactoryPrx.checkedCast(f)
            for i in range (0, number):
                spammers.append(fact.create(''))

        for sp in spammers:
            df.modify(FIPA.DfAgentDescription(sp, [FIPA.ServiceDescription('spammer')]))

    def infoUserSpammerPairs (self, comm):
        query = IceGrid.QueryPrx.checkedCast(comm.stringToProxy('FIPA/Query'))
        df = FIPA.DFPrx.checkedCast(query.findObjectById(comm.stringToIdentity('DF')))

        users = len(df.search(FIPA.DfAgentDescription(None, [FIPA.ServiceDescription('user')])))
        spammers = len(df.search(FIPA.DfAgentDescription(None, [FIPA.ServiceDescription('spammer')])))

        print 'User agents: ' + str(users)
        print 'Spammer agents: ' + str(spammers)

    def beginTest (self, messages, useMTS, messageSizeInBytes, comm):

        query = IceGrid.QueryPrx.checkedCast(comm.stringToProxy('FIPA/Query'))
        df = FIPA.DFPrx.checkedCast(query.findObjectById(comm.stringToIdentity('DF')))
        spammers = df.search(FIPA.DfAgentDescription(None, [FIPA.ServiceDescription('spammer')]))
        sniffer = Test.SnifferPrx.checkedCast(comm.stringToProxy('Sniffer'))

        sniffer.setConfiguration(useMTS, messages, messageSizeInBytes)

        cb = AMI_Model_beginTestI()
        for agent in spammers:
            sp = Test.SpammerPrx.uncheckedCast(agent)
            sp.beginTest_async(cb, useMTS, messages, messageSizeInBytes)

        comm.waitForShutdown()

    def clean(self, comm):
        query = IceGrid.QueryPrx.checkedCast(comm.stringToProxy('FIPA/Query'))
        df = FIPA.DFPrx.checkedCast(query.findObjectById(comm.stringToIdentity('DF')))
        
        spammers = df.search(FIPA.DfAgentDescription(None, [FIPA.ServiceDescription('spammer')]))
        users = df.search(FIPA.DfAgentDescription(None, [FIPA.ServiceDescription('user')]))

        for agent in spammers:
            agent.destroy()

        for agent in users:
            agent.destroy()

Client().main(sys.argv, os.path.join(HOME, 'test', 'config.client'))
